import os
from mistralai import Mistral

# --- 配置 ---
API_KEY = "Kz2yZgyEQxkreQRagET0rtz15l8smbRT"

# Mistral AI 所有已知模型列表
ALL_MODELS = [
    # 开源模型
    "open-mistral-7b",
    "open-mixtral-8x7b",
    "open-mixtral-8x22b",
    "open-mistral-nemo",
    "open-codestral-mamba",
    
    # 优化模型
    "mistral-small-latest",
    "mistral-small-2409",
    "mistral-small-2506",
    "mistral-medium-latest",
    "mistral-medium-2312",
    "mistral-medium-3",
    "mistral-large-latest",
    "mistral-large-2402",
    "mistral-large-2407",
    "mistral-large-2411",
    
    # 视觉模型
    "pixtral-12b-2409",
    "pixtral-large-latest",
    
    # 专用模型
    "codestral-latest",
    "codestral-2405",
    "codestral-mamba-latest",
    "mistral-embed",
    "mistral-moderation-latest",
    
    # 其他可能的模型名
    "mistral-tiny",
    "mistral-small",
    "mistral-medium",
]

def get_available_models_from_api():
    """从 API 获取可用模型列表"""
    print("="*60)
    print("从 Mistral API 获取可用模型列表")
    print("="*60)
    
    try:
        client = Mistral(api_key=API_KEY)
        
        # 调用 list_models API
        models_response = client.models.list()
        
        print(f"\n✓ 成功获取模型列表!")
        print(f"总共 {len(models_response.data)} 个模型可用\n")
        
        available_models = []
        print(f"{'模型ID':<40} {'创建日期':<15} {'拥有者'}")
        print("-"*80)
        
        for model in models_response.data:
            model_id = model.id
            created = model.created if hasattr(model, 'created') else 'N/A'
            owned_by = model.owned_by if hasattr(model, 'owned_by') else 'N/A'
            
            print(f"{model_id:<40} {str(created):<15} {owned_by}")
            available_models.append(model_id)
        
        return available_models
        
    except Exception as e:
        print(f"✗ 获取模型列表失败: {e}")
        return []

def test_model_with_simple_query(model_name):
    """用简单的查询测试模型是否可用"""
    try:
        client = Mistral(api_key=API_KEY)
        
        # 发送一个简单的测试消息
        response = client.chat.complete(
            model=model_name,
            messages=[
                {
                    "role": "user",
                    "content": "Hello, say 'OK' if you can hear me."
                }
            ],
            max_tokens=10
        )
        
        reply = response.choices[0].message.content
        return True, reply
        
    except Exception as e:
        error_msg = str(e)
        return False, error_msg

def test_all_models(model_list):
    """测试所有模型的可用性"""
    print("\n" + "="*60)
    print("测试各模型的可用性")
    print("="*60)
    
    results = {
        "available": [],
        "unavailable": []
    }
    
    for i, model in enumerate(model_list, 1):
        print(f"\n[{i}/{len(model_list)}] 测试模型: {model}")
        print("-"*60)
        
        success, response = test_model_with_simple_query(model)
        
        if success:
            print(f"✓ 可用")
            print(f"  回复: {response[:100]}")  # 只显示前100个字符
            results["available"].append(model)
        else:
            print(f"✗ 不可用")
            # 简化错误信息显示
            if "not found" in response.lower() or "does not exist" in response.lower():
                print(f"  原因: 模型不存在")
            elif "permission" in response.lower() or "authorized" in response.lower():
                print(f"  原因: 没有权限访问")
            elif "quota" in response.lower() or "limit" in response.lower():
                print(f"  原因: 配额限制")
            else:
                print(f"  原因: {response[:150]}")
            results["unavailable"].append(model)
    
    return results

def print_summary(results):
    """打印测试结果汇总"""
    print("\n" + "="*60)
    print("测试结果汇总")
    print("="*60)
    
    print(f"\n✓ 可用模型 ({len(results['available'])} 个):")
    print("-"*60)
    if results['available']:
        for model in results['available']:
            print(f"  • {model}")
    else:
        print("  (无)")
    
    print(f"\n✗ 不可用模型 ({len(results['unavailable'])} 个):")
    print("-"*60)
    if results['unavailable']:
        for model in results['unavailable']:
            print(f"  • {model}")
    else:
        print("  (无)")
    
    # 计算成功率
    total = len(results['available']) + len(results['unavailable'])
    if total > 0:
        success_rate = (len(results['available']) / total) * 100
        print(f"\n成功率: {success_rate:.1f}% ({len(results['available'])}/{total})")

# --- 主程序 ---
if __name__ == "__main__":
    print("="*60)
    print("Mistral AI 模型可用性测试工具")
    print("="*60)
    print()
    
    # 选项1: 从API获取官方模型列表
    print("选项1: 从 API 获取官方可用模型列表")
    api_models = get_available_models_from_api()
    
    # 选项2: 测试预定义的模型列表
    print("\n" + "="*60)
    print("选项2: 测试预定义的模型列表")
    print("="*60)
    print(f"将测试 {len(ALL_MODELS)} 个模型...")
    
    input("\n按 Enter 键开始测试预定义模型列表 (或 Ctrl+C 取消)...")
    
    results = test_all_models(ALL_MODELS)
    
    # 打印汇总
    print_summary(results)
    
    # 额外: 比较 API 列表和测试结果
    if api_models:
        print("\n" + "="*60)
        print("额外信息: API 列表中但未在预定义列表中的模型")
        print("="*60)
        
        missing_models = set(api_models) - set(ALL_MODELS)
        if missing_models:
            print("\n以下模型在 API 中可用,但未包含在测试列表中:")
            for model in sorted(missing_models):
                print(f"  • {model}")
        else:
            print("\n所有 API 模型都已包含在测试列表中")
    
    print("\n" + "="*60)
    print("测试完成!")
    print("="*60)